
######################################
#
#   cpu和gpu(cuda）运行时间对比
#
#####################################

import torch
import time


def gpu_cpu_run_time(a,b):
    """
    cpu和gpu运行时间对比：
    :return: none
    """
    t0 = time.time()
    c = torch.matmul(a, b)  # 在cpu用上进行矩阵乘法
    t1 = time.time()
    print(f"运行时间：{t1 - t0}，相乘后的矩阵形状：{c.shape},a在{a.device}上,b在{b.device}上, c在{c.device}上")  # 输出矩阵乘法在哪执行（cpu还是gpu）和执行浪费的时间

    """
    c.norm是求矩阵的Frobenius范数（根号下行||列中每个值平方和，得到的是一个向量），默认是行的范数，和数值分析中二范数求法（求AT{A的转置}*A的最大特征值）不一样
    """


if __name__ == '__main__':

    print(f"torch.__version__:{ torch.__version__ }")           # 输出torch的版本号
    print(f"can torch use gpu?:{ torch.cuda.is_available() }")  # 输出torch是否可以使用gpu(cuda)

    a = torch.randn(10000, 1500)  # 使用torch生成地址在cpu上的10000*1500行的矩阵
    b = torch.randn(1500, 20000)

    gpu_cpu_run_time(a,b)        #使用矩阵乘法来对比cpu运行和GPU运行的时间

    device=torch.device("cuda:0")     # 引出cuda地址
    a = a.to(device)                  # 将a的地址引入cuda（gpu）中
    b = b.to(device)

    gpu_cpu_run_time(a, b)           # 这里会初始化cuda环境，时间可能会增加
    torch.cuda.empty_cache()         # 释放显存

    a= torch.randn((10000,1500))  # 不释放显存，1022在我电脑再增加一列都会内存溢出
    a = a.to(device)
    b= torch.randn((1500,20000))
    b=b.to(device)
    gpu_cpu_run_time(a, b)